/*
*    l3qd - Light, light, lightweight queue
*    Copyright (C) 2024  Marcus Pedersén marcus@marcux.org
*
*    This program is free software: you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation, either version 3 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"gorm.io/gorm"
	"io"
	"log/slog"
	"net/http"
	"notabug.org/marcux/l3q/internal/common"
	"notabug.org/marcux/l3q/internal/daemon"
	"notabug.org/marcux/l3q/internal/daemon/sql"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Represent task
// that should be
// started on specific
// node.
// node, job and
// task structs contain values
// that should be updated.
type Run struct {
	node  common.Node
	tasks []common.Task
}

// Checks all nodes
// and update node status
func nodeStatus(db *gorm.DB) error {
	nodes, err := sql.GetNode(db, -1, false, common.NodeState(-1))
	if err != nil {
		return fmt.Errorf("Worker: Failed to get all nodes from database: %s", err)
	}

	for _, n := range nodes {
		jNode := common.Node{}

		resp, err := http.Get(fmt.Sprintf("https://%s/api/v1/daemon/status", n.Ip))
		if err != nil {
			resp, err = http.Get(fmt.Sprintf("https://%s/api/v1/daemon/status", n.Hostname))
			if err != nil {
				sErr := sql.SetNodeOffline(db, n.Id)
				if sErr != nil {
					slog.Error(fmt.Sprintf("Worker: %s", sErr))
				}
				continue
			}
		}

		defer resp.Body.Close()

		body, bErr := io.ReadAll(resp.Body)
		if bErr != nil {
			slog.Error(fmt.Sprintf("Worker: Failed to read body response from node with id: %d, %s", n.Id, bErr))
			continue
		}

		err = json.Unmarshal(body, &jNode)
		if err != nil {
			slog.Error(fmt.Sprintf("Worker: Failed to parse response from node with id: %d, %s", n.Id, err))
			continue
		}

		if n.State == common.NodeStateSoftOnline || n.State == common.NodeStateRegister || n.State == common.NodeStateOffline {
			jNode.State = common.NodeStateOnline
		}

		err = sql.UpdateNode(db, n.Id, jNode)
		if err != nil {
			slog.Error(fmt.Sprintf("Worker: %s", err))
		}
	}

	return nil
}

// Check node that are in
// MaintenanceDrain state
// if no tasks are running
// on node, status will
// be changed to maintenance
func nodeStatusDrain(db *gorm.DB) {
	nodes, err := sql.GetNode(db, -1, true, common.NodeStateMaintenanceDrain)
	if err != nil {
		slog.Error(fmt.Sprintf("Worker: Failed to get all nodes with state NodeStateMaintenanceDrain: %s", err))
	}

	for _, n := range nodes {
		if len(n.Tasks) == 0 {
			mNode := common.Node{State: common.NodeStateMaintenance}
			err := sql.UpdateNode(db, n.Id, mNode)
			if err != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to update status to maintenance onnode with id: %d, %s", n.Id, err))
			}
		}
	}

}

// Checks all tasks on running jobs.
// If all tasks has terminated
// successfully or with error
// job state will bbe updated.
func updateRunningJobs(db *gorm.DB) {
	running, err := sql.GetJob(db, -1, true, common.JobStateRunning, false, false, 0)
	if err != nil {
		slog.Error(fmt.Sprintf("Worker: Failed to get all running jobs: %s", err))
		return
	}

	for _, r := range running {
		noTasksRunning := 0
		noTasksFinished := 0
		noTasksErr := 0

		for _, t := range r.Tasks {
			if len(t.StartTime) > 0 && len(t.EndTime) == 0 && t.ExitCode == 0 {
				noTasksRunning += 1
			} else if len(t.StartTime) > 0 && len(t.EndTime) > 0 && t.ExitCode == 0 && t.SubState == "exited" {
				noTasksFinished += 1
			} else if len(t.StartTime) > 0 && len(t.EndTime) > 0 && (t.ExitCode != 0 || t.SubState == "failed") {
				noTasksErr += 1
			}
		}

		var updJob common.Job
		update := false

		if r.NoTasksRunning != noTasksRunning || r.NoTasksFinished != noTasksFinished || r.NoTasksErr != noTasksErr {
			updJob.NoTasksRunning = noTasksRunning
			updJob.NoTasksFinished = noTasksFinished
			updJob.NoTasksErr = noTasksErr
			update = true
		}

		if len(r.Tasks) == (noTasksFinished+noTasksErr) && noTasksRunning == 0 {
			update = true
			if noTasksErr == 0 {
				updJob.State = common.JobStateTerminated
			} else {
				updJob.State = common.JobStateError
			}
		}

		if update {
			if err := sql.UpdateJob(db, r.Id, updJob); err != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to update status: %s and tast status, on jobid: %d, %s", string(updJob.State), r.Id, err))
			}
		}
	}
}

// Update job that has status
// depend. If all depending jobs
// has terminated, job will be
// added to queue.
func updateDependJobs(db *gorm.DB) {
	dJobs, err := sql.GetJob(db, -1, false, common.JobStateDepend, false, false, 0)
	if err != nil {
		slog.Error(fmt.Sprintf("Worker: Not updating depending jobs: %s", err))
		return
	}

	for _, dep := range dJobs {
		allDepTerm := true
		depFail := false
		depends := strings.Split(strings.TrimSpace(dep.Depend), ",")
		for _, dStr := range depends {
			d, iErr := strconv.Atoi(strings.TrimSpace(dStr))
			if iErr != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to parse depend int: %s, for job id:%d, %s", dStr, dep.Id, iErr))
				allDepTerm = false
				break
			}

			dj, djErr := sql.GetJob(db, d, false, common.JobState(-1), false, false, 0)
			if djErr != nil {
				slog.Error(fmt.Sprintf("Workert: Failed to get depending job with id: %d, %s", d, djErr))
				allDepTerm = false
				break
			}

			if len(dj) != 1 {
				slog.Error(fmt.Sprintf("Worker: Wrong number of job returned: %d, when getting depending jobid: %d frim database", len(dj), d))
				allDepTerm = false
				break
			}

			if dj[0].State == common.JobStateQueued || dj[0].State == common.JobStateRunning || dj[0].State == common.JobStateDepend {
				allDepTerm = false
				break
			} else if dj[0].State != common.JobStateTerminated {
				allDepTerm = false
				depFail = true
				break
			}
		}

		if allDepTerm {
			jobState := common.Job{State: common.JobStateQueued}
			if err := sql.UpdateJob(db, dep.Id, jobState); err != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to change jobid: %d state to: JobStateQueued, %s", dep.Id, err))
			}
		} else if depFail {
			jobState := common.Job{State: common.JobStateDependError}
			if err := sql.UpdateJob(db, dep.Id, jobState); err != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to change jobid: %d state to: JobStateDependError, %s", dep.Id, err))
			}
		}
	}
}

// Sort nodes on available
// cpus
func sortNodes(nodes []common.Node) {
	slices.SortFunc(nodes, func(a, b common.Node) int {
		if a.CpusAlloc < b.CpusAlloc {
			return -1
		} else if a.CpusAlloc > b.CpusAlloc {
			return 1
		} else {
			return 0
		}
	})
}

// Runs the default queue
// or the fill queue
// and returns tasks that
// should be started.
// If defaultQ is true the
// default queue will run
// else the fill queue.
// Returns a map of nodeid
// with tasks to start and
// slice of jobs that should
// be updated.
// Default queue:
// Add all jobs in turn with
// the oldest job first.
// Returns when first job
// do not fit on one of the
// nodes.
// Fill queue:
// Add all jobs in turn with
// the oldest job first.
// Returns when as many jobs
// has been able to fit.
// Will skip job and try to
// fit next job until no more
// jobs will fit.
func queueDefaultFill(jobs []common.Job, nodes []common.Node, defaultQ bool) (map[int]Run, []common.Job) {
	nodesRun := make(map[int]Run)
	updJobs := []common.Job{}

	sortNodes(nodes)

	for _, j := range jobs {
		paraRun := make(map[int]Run)

		for i, n := range nodes {
			unallocCpus := n.CpusTotal - n.CpusAlloc
			if j.CoresAlloc <= unallocCpus {
				now := time.Now()
				nStr := ""

				if len(j.Nodes) == 0 {
					nStr = n.Name
				} else {
					if !strings.Contains(j.Nodes, n.Name) {
						nStr = fmt.Sprintf("%s,%s", j.Nodes, n.Name)
					}
				}

				updTask := common.Task{
					Id:     j.Tasks[0].Id,
					NodeID: uint(n.Id),
				}

				if j.Type == common.JobTypeSequence {
					updJob := common.Job{
						Id:             j.Id,
						Type:           j.Type,
						StartDate:      now.Format("2006-01-02 15:04:05"),
						NoTasksRunning: j.NoTasksRunning + 1,
						Nodes:          nStr,
						State:          common.JobStateRunning,
					}

					if _, ok := nodesRun[n.Id]; ok == true {
						runN := nodesRun[n.Id].node
						runN.CpusAlloc = runN.CpusAlloc + j.CoresAlloc
						runT := nodesRun[n.Id].tasks
						runT = append(runT, updTask)
						nodesRun[n.Id] = Run{node: runN, tasks: runT}
					} else {
						r := Run{}
						r.node = common.Node{
							Id:        n.Id,
							CpusAlloc: n.CpusAlloc + j.CoresAlloc,
							Ip:        n.Ip,
							Hostname:  n.Hostname,
							Port:      n.Port,
							Key:       n.Key,
							DaemonKey: n.DaemonKey,
						}

						r.tasks = append(r.tasks, updTask)
						nodesRun[n.Id] = r
					}

					updJobs = append(updJobs, updJob)
					n.CpusAlloc = n.CpusAlloc + j.CoresAlloc
					nodes[i] = n
					break
				} else {
					if _, ok := paraRun[n.Id]; ok == true {
						runN := paraRun[n.Id].node
						runN.CpusAlloc = runN.CpusAlloc + j.CoresAlloc
						runT := paraRun[n.Id].tasks
						runT = append(runT, updTask)
						paraRun[n.Id] = Run{node: runN, tasks: runT}
					} else {
						r := Run{}
						r.node = common.Node{
							Id:        n.Id,
							CpusAlloc: n.CpusAlloc + j.CoresAlloc,
							Ip:        n.Ip,
							Hostname:  n.Hostname,
							Port:      n.Port,
							Key:       n.Key,
							DaemonKey: n.DaemonKey,
						}

						if _, ok := nodesRun[n.Id]; ok == true {
							r.tasks = nodesRun[n.Id].tasks
						}

						r.tasks = append(r.tasks, updTask)
						paraRun[n.Id] = r
					}

					if len(j.StartDate) == 0 {
						j.StartDate = now.Format("2006-01-02 15:04:05")
					}

					j.NoTasksRunning = j.NoTasksRunning + 1
					j.State = common.JobStateRunning
					j.Nodes = nStr

					if len(j.Tasks) == 1 {
						updJob := common.Job{
							Id:             j.Id,
							Type:           j.Type,
							StartDate:      j.StartDate,
							NoTasksRunning: j.NoTasksRunning,
							Nodes:          j.Nodes,
							State:          j.State,
						}

						for id, r := range paraRun {
							nodesRun[id] = r
						}

						updJobs = append(updJobs, updJob)
						n.CpusAlloc = n.CpusAlloc + j.CoresAlloc
						nodes[i] = n
						break
					}

					j.Tasks = j.Tasks[1:]
					n.CpusAlloc = n.CpusAlloc + j.CoresAlloc
					nodes[i] = n
				}
			} else {
				if defaultQ {
					return nodesRun, updJobs
				} else {
					break
				}
			}
		}
		sortNodes(nodes)
	}

	return nodesRun, updJobs
}

// Connects to nodes
// and send jobs in map
// to specified node
func startJobs(run map[int]Run, updateJobs []common.Job) []common.Job {
	failedJobs := []int{}

	for _, r := range run {
		// check jobid list on tasks that should
		// be removed from r.tasks
		for _, f := range failedJobs {
			for {
				i := slices.IndexFunc(r.tasks, func(t common.Task) bool {
					return int(t.JobID) == f
				})

				if i != -1 {
					slices.Delete(r.tasks, i, i+1)
				} else {
					break
				}
			}
		}

		url := fmt.Sprintf("http://%s:%d/api/v1/daemon/task", r.node.Hostname, r.node.Port)
		ipUrl := fmt.Sprintf("http://%s:%d/api/v1/daemon/task", r.node.Ip, r.node.Port)
		jBody, err := json.Marshal(r.tasks)
		if err != nil {
			slog.Error(fmt.Sprintf("Worker: Failed to convert task list to json for nodeid: %d, %s", r.node.Id, err))
			// update job
			for _, t := range r.tasks {
				j := slices.IndexFunc(updateJobs, func(j common.Job) bool {
					return j.Id == int(t.JobID)
				})

				if j != -1 {
					job := updateJobs[j]
					if job.Type == common.JobTypeParallel {
						if !slices.Contains(failedJobs, job.Id) {
							failedJobs = append(failedJobs, job.Id)
						}
					}

					// delete job from updateJobs
					if j == len(updateJobs)-1 {
						updateJobs = updateJobs[:j]
					} else if j == len(updateJobs)-2 {
						updateJobs = append(updateJobs[:j], updateJobs[j+1])
					} else {
						updateJobs = append(updateJobs[:j], updateJobs[j+1:len(updateJobs)-1]...)
					}
				}
			}
			continue
		}

		byteJbody := bytes.NewBuffer(jBody)
		resp, pErr := http.Post(url, "application/json", byteJbody)
		if pErr != nil {
			resp, pErr = http.Post(ipUrl, "application/json", byteJbody)
			if pErr != nil {
				slog.Error(fmt.Sprintf("Worker: Failed to send job to nodeid: %d, url: %s, %s, %s", r.node.Id, ipUrl, resp.Status, pErr))
				// update job
				for _, t := range r.tasks {
					j := slices.IndexFunc(updateJobs, func(j common.Job) bool {
						return j.Id == int(t.JobID)
					})

					if j != -1 {
						job := updateJobs[j]
						if job.Type == common.JobTypeParallel {
							if !slices.Contains(failedJobs, job.Id) {
								failedJobs = append(failedJobs, job.Id)
							}
						}

						// delete job from updateJobs
						if j == len(updateJobs)-1 {
							updateJobs = updateJobs[:j]
						} else if j == len(updateJobs)-2 {
							updateJobs = append(updateJobs[:j], updateJobs[j+1])
						} else {
							updateJobs = append(updateJobs[:j], updateJobs[j+1:len(updateJobs)-1]...)
						}
					}
				}
				continue
			}
		}

		if resp.StatusCode != 200 {
			slog.Error(fmt.Sprintf("Worker: Failed to send job to nodeid: %d, url: %s, wrong status code in response: %s", r.node.Id, url, resp.Status))
			// update job
			for _, t := range r.tasks {
				j := slices.IndexFunc(updateJobs, func(j common.Job) bool {
					return j.Id == int(t.JobID)
				})

				if j != -1 {
					job := updateJobs[j]
					if job.Type == common.JobTypeParallel {
						if !slices.Contains(failedJobs, job.Id) {
							failedJobs = append(failedJobs, job.Id)
						}
					}

					// delete job from updateJobs
					if j == len(updateJobs)-1 {
						updateJobs = updateJobs[:j]
					} else if j == len(updateJobs)-2 {
						updateJobs = append(updateJobs[:j], updateJobs[j+1])
					} else {
						updateJobs = append(updateJobs[:j], updateJobs[j+1:len(updateJobs)-1]...)
					}
				}
			}
		}
	}

	return updateJobs
}

// Get queued data from database
// and runs specified queue.
func updateQueue(db *gorm.DB, queue string) {
	qJobs, err := sql.GetJob(db, -1, true, common.JobStateQueued, false, false, 0)
	if err != nil {
		slog.Error(fmt.Sprintf("Worker: Failed to get all queued jobs, not running queue: %s", err))
		return
	}

	onlineNodes, nErr := sql.GetNode(db, -1, true, common.NodeStateOnline)
	if nErr != nil {
		slog.Error(fmt.Sprintf("Worker: Failed to get all online nodes, not running queue: %s", nErr))
		return
	}

	var run map[int]Run
	var updateJobs []common.Job

	switch queue {
	case "default":
		run, updateJobs = queueDefaultFill(qJobs, onlineNodes, true)
	case "fill":
		run, updateJobs = queueDefaultFill(qJobs, onlineNodes, false)
	default:
		slog.Error("Worker: Failed to run queue: %s, queue not defined", queue)
	}

	updateJobs = startJobs(run, updateJobs)

	for _, j := range updateJobs {
		err := sql.UpdateJob(db, j.Id, j)
		if err != nil {
			slog.Error(fmt.Sprintf("Worker: Failed to update just started jobs: %s", err))
		}
	}
}

// Starts and run the worker
// process. Checks channel
// and if true the worker
// process terminates
// pollInterval is time in
// seconds between poll of
// node status and queue update
func RunWorker(conf daemon.Config, end chan bool, wait *sync.WaitGroup) {
	defer wait.Done()

	slog.Info("Worker: Process started")

	db, err := sql.InitDb(conf)
	if err != nil {
		slog.Error(fmt.Sprintf("Worker: Failed to open database: %s", err))
		return
	}

	noTicks := 0
	sec := time.Tick(time.Second)

	for _ = range sec {
		noTicks++
		if len(end) > 0 {
			e, ok := <-end

			if e && ok {
				break
			} else if !ok {
				break
			}
		} else if noTicks >= conf.PollIntervall {
			if err := nodeStatus(db); err != nil {
				slog.Error(fmt.Sprintf("Worker: %s", err))
			}

			nodeStatusDrain(db)
			updateRunningJobs(db)
			updateDependJobs(db)
			updateQueue(db, conf.Queue)

			noTicks = 0
		}
	}

	slog.Info("Worker: Process terminated")
}
